import asyncio
import os
from mcp import ClientSession, StdioServerParameters, types
from mcp.client.stdio import stdio_client
from pydantic import AnyUrl


async def main():
    base_dir = os.path.dirname(os.path.abspath(__file__))
    server_path = os.path.join(base_dir, "server.py")
    server_params = StdioServerParameters(command="python", args=[server_path])
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            prompts = await session.list_prompts()
            print(f"[Prompts]", [p.name for p in prompts.prompts])

            resource_templates = await session.list_resource_templates()
            print(
                "[Resource Template]",
                [t.uriTemplate for t in resource_templates.resourceTemplates],
            )
            language = None
            style = None
            if prompts.prompts:
                prompt_name = prompts.prompts[0].name
                print(f"Prompt补全示例:{prompt_name}")
                language_completion = await session.complete(
                    ref=types.PromptReference(type="ref/prompt", name=prompt_name),
                    argument={"name": "language", "value": "py"},
                )
                print(f"language_completion: {language_completion.completion.values}")
                language = language_completion.completion.values[0]
                style_completion = await session.complete(
                    ref=types.PromptReference(type="ref/prompt", name=prompt_name),
                    argument={"name": "style", "value": "func"},
                )
                print(f"style_completion: {style_completion.completion.values}")
                style = style_completion.completion.values[0]
                result = await session.get_prompt(
                    name="code_style_guide",
                    arguments={"language": language, "style": style},
                )
                print("code_style_guide的提示词是", result.messages[0].content.text)

            if resource_templates.resourceTemplates:
                resource_template = resource_templates.resourceTemplates[0]
                print(f"资源模板补全示例 {resource_template.uriTemplate}")
                owner_completion = await session.complete(
                    ref=types.ResourceTemplateReference(
                        type="ref/resource", uri=resource_template.uriTemplate
                    ),
                    argument={"name": "owner", "value": "mod"},
                )
                print(f"owner_completion: {owner_completion.completion.values}")
                owner = owner_completion.completion.values[0]
                repo_completion = await session.complete(
                    ref=types.ResourceTemplateReference(
                        type="ref/resource", uri=resource_template.uriTemplate
                    ),
                    argument={"name": "repo", "value": "py"},
                    context_arguments={"owner": "modelcontextprotocol"},
                )
                print(f"repo_completion: {repo_completion.completion.values}")
                repo = repo_completion.completion.values[0]
                result = await session.read_resource(AnyUrl(f"github://{owner}/{repo}"))
                print(f"github://{owner}/{repo}是", result.contents[0].text)


if __name__ == "__main__":
    asyncio.run(main())
